#include <thrust/copy.h>
#include <cusp/precond/spai.h>
#include <cusp/krylov/gmres.h>
#include <cusp/gallery/poisson.h>
#include <cusp/gallery/eye.h>
#include <cusp/csr_matrix.h>

#include <cusp/frobenius_norm.h>
#include <cusp/print.h>
#include <cusp/multiply.h>
#include <cusp/elementwise.h>
#include <cusp/csr_matrix.h>
#include <cusp/graph/neighborhood_extension.h>

#include <iostream>

#define NEIGH_LEV 4

// Number of levels of neighbors to be considered for considered

int main(int argc, char** argv)
{
    typedef int                 IndexType;
    typedef float		ValueType;
    typedef cusp::device_memory MemorySpace;

    // create an empty sparse matrix structure
    cusp::csr_matrix<IndexType, ValueType, cusp::host_memory> A_csr;

    if (argc < 2)
    {
        cusp::gallery::poisson5pt(A_csr, 512, 512);
    }
    else
    {
        // read matrix from file
        cusp::io::read_matrix_market_file(A_csr, std::string(argv[1]));
    }

    std::cout << "size " << A_csr.num_rows << " x " << A_csr.num_cols << std::endl;

    cusp::csr_matrix<IndexType, ValueType, MemorySpace> d_A(A_csr);
    cusp::csr_matrix<IndexType, ValueType, MemorySpace> I;
    cusp::gallery::eye(I, d_A.num_rows, d_A.num_cols);
    cusp::csr_matrix<int,double,cusp::device_memory> P;
    cusp::subtract(d_A, I, P);
    std::cout << "Fnorm(A-I)\t" << cusp::frobenius_norm(P) << "\n";

//    cusp::print(SP_csr);

    // solve SPAI preconditioner
    for (IndexType lev=0; lev< NEIGH_LEV; lev++)
    {
        // cusp::csr_matrix<IndexType, ValueType, cusp::host_memory> poi_csr; // Copy constructor
        cusp::csr_matrix<IndexType, ValueType, cusp::host_memory> SP_csr(cusp::graph::neighborhood_extension(A_csr,lev)); // Copy constructor

        cusp::csr_matrix<IndexType, ValueType, MemorySpace> d_SP(SP_csr);
        std::cout << "\nSolving with SPAI preconditioner with neighbor level " << lev << " non-zeros: " << SP_csr.num_entries << std::endl;
        
        // setup preconditioner
        cusp::precond::spai<ValueType, MemorySpace> M(d_A, d_SP);

        std::cout << "Completed preconditioner \n";

        // report norms
        cusp::multiply(d_A, M.matrix, P);
        cusp::subtract(P, I, P);
        std::cout << "Fnorm(AM-I)\t" << cusp::frobenius_norm(P) << std::endl;
    }

    return 0;
}

